package com.tsx.ai.service.impl;

import com.tsx.ai.service.ChatService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;

@RequiredArgsConstructor
@Slf4j
@Service
public class ChatServiceImpl implements ChatService {
    private final ChatClient chatClient;
    private final VectorStore vectorStore;

    @Override
    public String chat(String sessionId, String question) {
        // 创建搜索请求，用于搜索相关文档
        SearchRequest searchRequest = SearchRequest.builder()
                .query(question) // 设置查询条件
                .topK(3) // 设置最多返回的文档数量
                .build();
        // 创建一个QuestionAnswerAdvisor，用于处理搜索结果
        QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(searchRequest)
                .build();
        String content = chatClient.prompt(question)
                .system(promptSystemSpec -> promptSystemSpec.param("now", LocalDateTime.now()))
                .advisors(advisorSpec -> advisorSpec
                        .param(ChatMemory.CONVERSATION_ID, sessionId)
                        .advisors(questionAnswerAdvisor))
                .call()
                .content();
        log.info("question :{},content :{}", question, content);
        return content;
    }

    @Override
    public Flux<String> chatStream(String sessionId, String question) {
        return chatClient.prompt(question)
                .stream()
                .content()
                // 每次收到响应内容就运行
                .doOnNext(content -> log.info("question :{},content :{}", question, content))
                // 在流结束后添加结束标记
                .concatWith(Flux.just("[END]"));
    }
}
